
from evaluator.rag_agent import RagAgent
import requests
import json


class DifyRagAgent(RagAgent):
    def __init__(self, api_url: str, api_key: str):
        self.api_url = api_url
        self.api_key = api_key

    def generate_answer(self, question: str, contexts: list[str]) -> str:
        # 设置请求头
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }
        payload = {
            "inputs": {},
            "query": question,
            "response_mode": "blocking",
            "conversation_id": None,
            "user": "abc-123"
        }
        
        response = requests.post(self.api_url + 'chat-messages', headers=headers, data=json.dumps(payload))
        
        if response.status_code == 200:
            return response.json().get('answer')
        else:
            print(f"Error: Received status code {response.status_code}")
            print(f"Response body: {response.text}")
            return None
